import json
import requests
from dashscope import Generation

# 1. 定义网络搜索函数的schema
web_search_schema = {
    "name": "web_search", # name 指定函数名称
    "description": "搜索互联网上的信息，返回相关的搜索结果，支持知识图谱和本地商户信息",
     # parameters 定义函数的入参格式
    "parameters": { 
        "type": "object",
        "properties": {
            "q": {
                "type": "string",
                "description": "搜索查询关键词，例如'人工智能最新发展'或'餐厅名称'"
            },
            "engine": {
                "type": "string",
                "description": "搜索引擎选择，默认为google"
            },
            "location": {
                "type": "string",
                "description": "搜索位置，例如'New York'或'北京'，用于本地商户搜索"
            }
        },
        "required": ["q"]
    },
    # returns 定义函数的返回格式
    "returns": {
        "type": "object",
        "properties": {
            "knowledge_graph": {
                "type": "object",
                "description": "知识图谱信息，包含本地商户详细信息",
                "properties": {
                    "kgmid": {
                        "type": "string",
                        "description": "知识图谱ID"
                    },
                    "knowledge_graph_type": {
                        "type": "string",
                        "description": "知识图谱类型"
                    },
                    "title": {
                        "type": "string",
                        "description": "商户名称"
                    },
                    "type": {
                        "type": "string",
                        "description": "商户类型，如'法国餐厅'"
                    },
                    "rating": {
                        "type": "number",
                        "description": "评分，如4.6"
                    },
                    "reviews": {
                        "type": "integer",
                        "description": "评论数量"
                    },
                    "price": {
                        "type": "string",
                        "description": "价格等级，如'$$$$'"
                    },
                    "price_description": {
                        "type": "string",
                        "description": "价格描述，如'Very expensive'"
                    },
                    "website": {
                        "type": "string",
                        "description": "商户网站"
                    },
                    "directions": {
                        "type": "string",
                        "description": "导航链接"
                    },
                    "description": {
                        "type": "string",
                        "description": "商户描述"
                    },
                    "service_options": {
                        "type": "array",
                        "description": "服务选项，如['堂食','外卖']",
                        "items": {
                            "type": "string"
                        }
                    },
                    "address": {
                        "type": "string",
                        "description": "商户地址"
                    },
                    "menu": {
                        "type": "object",
                        "description": "菜单信息",
                        "properties": {
                            "text": {
                                "type": "string",
                                "description": "菜单文本"
                            },
                            "link": {
                                "type": "string",
                                "description": "菜单链接"
                            }
                        }
                    },
                    "phone": {
                        "type": "string",
                        "description": "联系电话"
                    },
                    "hours": {
                        "type": "string",
                        "description": "营业时间概述"
                    },
                    "open_hours": {
                        "type": "array",
                        "description": "详细营业时间",
                        "items": {
                            "type": "object",
                            "properties": {
                                "name": {
                                    "type": "string",
                                    "description": "星期几"
                                },
                                "value": {
                                    "type": "string",
                                    "description": "营业时间段"
                                }
                            }
                        }
                    },
                    "popular_times": {
                        "type": "object",
                        "description": "人流量信息",
                        "properties": {
                            "live": {
                                "type": "object",
                                "description": "实时人流量"
                            },
                            "chart": {
                                "type": "object",
                                "description": "人流量图表数据"
                            }
                        }
                    },
                    "web_reviews": {
                        "type": "array",
                        "description": "网站评论",
                        "items": {
                            "type": "object",
                            "properties": {
                                "title": {
                                    "type": "string",
                                    "description": "评论网站"
                                },
                                "link": {
                                    "type": "string",
                                    "description": "评论链接"
                                },
                                "rating": {
                                    "type": "number",
                                    "description": "评分"
                                },
                                "reviews": {
                                    "type": "integer",
                                    "description": "评论数量"
                                }
                            }
                        }
                    },
                    "user_reviews": {
                        "type": "array",
                        "description": "用户评论",
                        "items": {
                            "type": "object",
                            "properties": {
                                "review": {
                                    "type": "string",
                                    "description": "评论内容"
                                },
                                "rating": {
                                    "type": "number",
                                    "description": "评分"
                                },
                                "user": {
                                    "type": "object",
                                    "description": "用户信息"
                                }
                            }
                        }
                    },
                    "profiles": {
                        "type": "array",
                        "description": "社交媒体资料",
                        "items": {
                            "type": "object",
                            "properties": {
                                "name": {
                                    "type": "string",
                                    "description": "平台名称"
                                },
                                "link": {
                                    "type": "string",
                                    "description": "资料链接"
                                }
                            }
                        }
                    },
                    "people_also_search_for": {
                        "type": "array",
                        "description": "相关搜索",
                        "items": {
                            "type": "object",
                            "properties": {
                                "name": {
                                    "type": "string",
                                    "description": "相关商户名称"
                                },
                                "link": {
                                    "type": "string",
                                    "description": "相关商户链接"
                                }
                            }
                        }
                    },
                    "local_map": {
                        "type": "object",
                        "description": "地图信息",
                        "properties": {
                            "link": {
                                "type": "string",
                                "description": "地图链接"
                            },
                            "gps_coordinates": {
                                "type": "object",
                                "description": "GPS坐标",
                                "properties": {
                                    "latitude": {
                                        "type": "number",
                                        "description": "纬度"
                                    },
                                    "longitude": {
                                        "type": "number",
                                        "description": "经度"
                                    }
                                }
                            }
                        }
                    }
                }
            },
            "organic": {
                "type": "array",
                "description": "常规搜索结果列表",
                "items": {
                    "type": "object",
                    "properties": {
                        "title": {
                            "type": "string",
                            "description": "搜索结果标题"
                        },
                        "link": {
                            "type": "string",
                            "description": "搜索结果链接"
                        },
                        "snippet": {
                            "type": "string",
                            "description": "搜索结果摘要"
                        },
                        "domain": {
                            "type": "string",
                            "description": "搜索结果来源网站"
                        }
                    }
                }
            }
        }
    }
}

# 2. 实现网络搜索函数
def perform_web_search(query, num_results="5", language="zh"):
    """执行网络搜索，调用SearchAPI搜索引擎API"""
    # SearchAPI配置
    api_key = "J1ZQcvEd65XxpEPpWdbEdxAs"  # 请替换为你的SearchAPI密钥
    api_url_base = "https://www.searchapi.io/api/v1/search"
    
    # 构建请求参数
    params = {
        "engine": "baidu",  # 搜索引擎
        "api_key": api_key,  # 必填
        "q": query,  # 搜索查询 必填
        "location": language if language != "zh" else None  # 如果是非中文搜索，使用language作为location
    }
    
    # 移除None值
    params = {k: v for k, v in params.items() if v is not None}
    
    try:
        # 发送API请求
        print("请求SearchAPI入参值", params)
        response = requests.get(api_url_base, params=params)
        response.raise_for_status()  # 检查请求是否成功
        print("请求SearchAPI返回值", response)
        
        # 解析返回的JSON数据
        data = response.json()
        print("解析SearchAPI返回值", data)
        # 直接返回API响应数据，包含知识图谱和有机搜索结果
        # 如果有知识图谱数据，保留原始结构
        if "knowledge_graph" in data:
            print("检测到知识图谱数据")
            # 直接返回原始数据，保持与schema一致的结构
            result = data
        else:
            # 如果没有知识图谱，处理常规搜索结果
            organic_results = data.get("organic_results", [])
            print(f"获取到{len(organic_results)}条常规搜索结果")
            
            result = {
                "organic": organic_results
            }
        return result
            
    except requests.exceptions.RequestException as e:
        print(f"请求SearchAPI出错: {str(e)}")
        return {"total": "0", "results": []}
    except Exception as e:
        print(f"处理搜索数据出错: {str(e)}")
        return {"total": "0", "results": []}

# 3. 定义函数工具列表
function_tools = [
    {
        "type": "function",
        "function": web_search_schema # 函数工具的schema
    }
]

# 4. 函数调用处理器
def function_call_handler(function_call):
    """处理模型返回的函数调用请求"""
    arguments = json.loads(function_call.get('arguments', '{}'))
    
    query = arguments.get('q')
    location = arguments.get('location', '')
    language = 'en' if location else 'zh'  # 如果提供了location，默认使用英文
    
    if not query:
        print("搜索查询为空")
        return {"error": "搜索查询不能为空"}
    
    return perform_web_search(query, "5", language if not location else location)
    

# 调用Qwen模型并处理函数调用
def chat_with_function_calling(api_key, messages, tools=None):
    """使用DashScope SDK调用Qwen模型的函数调用功能"""
    try:
        # 第一次调用模型，可能会返回函数调用请求
        response = Generation.call(
            model='qwen-max',  # 或其他支持函数调用的Qwen模型
            api_key=api_key,
            messages=messages,
            tools=tools,  # 传入可调用的函数
            tool_choice="auto",  # 让模型自动决定是否调用函数
            result_format='message'  # 返回消息格式
        )
        
        # 检查是否有函数调用
        if response.output and response.output.choices:
            message = response.output.choices[0].message
            print("函数调用详情：", message)
            # 如果有函数调用
            if message.get('tool_calls'):
                tool_call = message['tool_calls'][0]
                function_call = tool_call.get('function')
                # 处理函数调用
                function_result = function_call_handler(function_call)
                # 将函数调用结果添加到消息历史
                messages.append({
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [{
                        "id": tool_call.get('id'),
                        "type": "function",
                        "function": function_call
                    }]
                })
                
                # 添加函数执行结果
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.get('id'),
                    "content": json.dumps(function_result, ensure_ascii=False)
                })
                
                # 再次调用模型，获取最终回复
                final_response = Generation.call(
                    model='qwen-max',
                    api_key=api_key,
                    messages=messages,
                    result_format='message'
                )
                
                if final_response.output and final_response.output.choices:
                    return final_response.output.choices[0].message.get('content')
                else:
                    return "模型未返回有效回复"
            else:
                # 没有函数调用，直接返回模型回复
                return message.get('content')
        else:
            return "模型未返回有效回复"
    except Exception as e:
        return f"调用出错: {str(e)}"

# 主函数
def main():
    # 请替换为你的DashScope API密钥
    api_key = "sk-***"
    # 用户消息
    messages = [
        {"role": "user", "content": "最近有哪些人工智能的重大突破？"}
    ]
    print("\n测试1: 搜索人工智能的最新突破")
    # 调用模型并处理函数调用
    response = chat_with_function_calling(api_key, messages, tools=function_tools)
    print("模型最终回复:")
    print(response)
    
    # 测试2: 英文搜索
    messages2 = [
        {"role": "user", "content": "What are the latest developments in quantum computing?"}
    ]
    print("\n测试2: 英文搜索量子计算的最新发展")
    response2 = chat_with_function_calling(api_key, messages2, tools=function_tools)
    print("模型最终回复:")
    print(response2)

# 运行示例
if __name__ == "__main__":
    main()